# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/5 10:54 
@Author  : ZhangShenao 
@File    : 1.使用原生API实现记忆功能.py 
@Desc    : 使用智谱AI的原生API,实现对话的记忆功能
"""
import os

import dotenv
from zhipuai import ZhipuAI

from summarize_buffer_memory import SummarizeBufferMemory

# 系统指令
SYSTEM_INSTRUMENT = '''你是一个具有记忆功能的AI助手，请根据结合用户的当前提问，和之前历史信息摘要和历史对话列表，生成准确的回答。你们的对话全部采用中文。
下面是之前的历史信息摘要：
<历史信息摘要开始>

{history_summarize}

<历史信息摘要结束>

下面是你们的历史聊天记录：
<历史聊天记录开始>

{chat_history_msg}

<历史聊天记录结束>
'''

# 调用的模型名称
MODEL_NAME = 'glm-4-air'

# 加载环境变量
dotenv.load_dotenv()

# 创建智普AI客户端
client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))

# 创建摘要缓冲记忆组件
memory = SummarizeBufferMemory(max_tokens=200)

# 从控制台接收输入,并传递给LLM
while True:
    query = input('Human: ')
    if 'bye' == query:
        print('bye bye~')
        break

    # 格式化历史摘要和历史聊天记录
    chat_history_msg = memory.format_chat_history()
    history_summarize = memory.get_history_summarize()

    # 流式调用LLM接口
    system_prompt = SYSTEM_INSTRUMENT.format(history_summarize=history_summarize, chat_history_msg=chat_history_msg)
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {'role': 'system', 'content': system_prompt},
            {"role": "user", "content": query},
        ],
        stream=True,
    )

    # 打印LLM的输出
    output = ''
    for chunk in response:
        content = chunk.choices[0].delta.content
        print(f'{content}', flush=True, end='')
        output += content
    print()

    # 保存历史聊天记录
    memory.save_chat_message(message=f'Human：{query}')
    memory.save_chat_message(message=f'AI：{output}')
